import time
from typing import Dict, Tuple, Optional
import gurobipy as gp
from gurobipy import GRB
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List

from mpl_toolkits.axisartist.angle_helper import select_step

from ChipSchduile.model.utils import print_schedule_results
from base import Task,Job,create_job
from ParamaGenerator import ParamterGenerator


class SemiconductorScheduler:
    def __init__(self, params, time_limit=600):
        """初始化模型参数并创建Gurobi模型，优先使用启发式求解"""

        self.params = params  # 模型参数字典
        self.model = gp.Model("Semiconductor_Scheduler")

        # 设置求解时间限制（秒）
        self.model.setParam('TimeLimit', time_limit)
        # 调整MIPGap，允许一定的最优性差距
        self.model.setParam('MIPGap', 0.01)
        # self.model.setParam('SolutionLimit', 1)  # 找到1个可行解就停止

        # 启发式算法增强参数
        # 提高启发式算法的迭代次数（默认值为10，增大可增强启发式搜索）
        self.model.setParam('Heuristics', 0.8)  # 取值范围[0,1]，越大越优先启发式
        # 减少精确算法（分支定界）的时间占比，优先启发式
        self.model.setParam('Cuts', 0)  # 关闭切割平面
        # 增强可行性启发式（快速寻找可行解），取值范围通常为 0~10（数值越大，启发式搜索强度越高）
        self.model.setParam('RINS', 10)
        # 允许更早启动启发式算法，优先在搜索初期启动启发式
        self.model.setParam('MIPFocus', 1)  # 优先启动启发式寻找可行解
        # 限制分支定界树的规模，迫使其依赖启发式
        # self.model.setParam('NodeLimit', 1000)  # 限制节点数量，超过后依赖启发式


        print("====================== begin 初始化决策变量 =========================")
        self.initialize_variables()  # 初始化决策变量
        print("====================== end 初始化决策变量 =========================")

        print("====================== begin 设置目标函数 =========================")
        self.set_objective()
        print("====================== end 设置目标函数 =========================")

        print("====================== begin 添加约束条件 =========================")
        self.add_constraints()
        print("====================== end 添加约束条件 =========================")

    def initialize_variables(self):
        """初始化所有决策变量"""

        J = self.params['J']  # 工序集合
        K = self.params['K']  # 机器集合
        R = self.params['R']  # 机械手集合
        PM = self.params['PM']  # 加工可行矩阵
        TM = self.params['TM']  # 运输可行矩阵
        Nj = self.params['Nj']

        # 决策变量定义
        self.Cmax = self.model.addVar(name="Cmax")  # 最大完工时间

        # 工序j的实际加工/实际运输时间
        self.Pt = self.model.addVars(J, name="Pt")
        self.Tt = self.model.addVars(J, name="Tt")

        # 工序j的加工/运输开始时刻
        self.SO = self.model.addVars(J, name="SO")  # 加工开始时间
        self.ST = self.model.addVars(J, name="ST")  # 运输开始时间

        # 加工任务分配变量 (1 if 工序j在机器k上加工)
        self.phi = self.model.addVars(J, K, vtype=GRB.BINARY, name="phi")

        # 运输任务分配变量 (1 if 工序j由机械手r运输)
        self.mu = self.model.addVars(J, R, vtype=GRB.BINARY, name="mu")

        # 机器k上工序j优先于j'的排序变量
        # self.y = self.model.addVars(J, J, K, vtype=GRB.BINARY, name="y")
        self.y = {}
        for j in J:
            for j_prime in J:
                if j != j_prime:
                    for k in K:
                        # 从加工可行字典中获取j,j_prime的加工机器
                        if PM.get((j, k), 0) == 1 and PM.get((j_prime, k), 0) == 1:
                            self.y[(j, j_prime, k)] = self.model.addVar(vtype=GRB.BINARY,
                                                                        name=f"y_{j}_{j_prime}_{k}")


        # 机械手r上工序j优先于j'的排序变量
        # self.z = self.model.addVars(J, J, R, vtype=GRB.BINARY, name="z")
        self.z = {}
        # for j in J:
        #     for j_prime in J:
        #         if j != j_prime:
        #             for r in R:
        #                 if TM.get((j, r), 0) == 1 and TM.get((j_prime, r), 0) == 1:
        #                     self.z[(j, j_prime, r)] = self.model.addVar(vtype=GRB.BINARY,
        #                                                                 name=f"z_{j}_{j_prime}_{r}")

        # 优化：仅对同一个工件上的工序创建这个决策变量
        # for group in Nj:
        #     for i in range(len(group) - 1):
        #         j_prev = group[i]  # 前序工序
        #         j_curr = group[i + 1]  # 后序工序（相邻）
        #         for r in self.params['R']:
        #             if TM.get((j_prev, r), 0) == 1 and TM.get((j_curr, r), 0) == 1:
        #                 # 仅为相邻工序创建z变量
        #                 self.z[(j_prev, j_curr, r)] = self.model.addVar(
        #                     vtype=GRB.BINARY, name=f"z_{j_prev}_{j_curr}_{r}")
        # 遍历所有任务对和机械手，提前初始化可能用到的z变量
        for r in R:
            for j in J:
                for j_prime in J:
                    if j != j_prime:
                        # 检查 (j, j_prime, r) 是否在self.z中，不在就添加
                        if (j, j_prime, r) not in self.z:
                            self.z[(j, j_prime, r)] = self.model.addVar(
                                vtype=GRB.BINARY,
                                name=f"z_{j}_{j_prime}_{r}"
                            )
                        # 检查 (j_prime, j, r) 是否在self.z中，不在就添加
                        if (j_prime, j, r) not in self.z:
                            self.z[(j_prime, j, r)] = self.model.addVar(
                                vtype=GRB.BINARY,
                                name=f"z_{j_prime}_{j}_{r}"
                            )

        self.model.update()

    def set_objective(self):
        """设置目标函数: 最小化最大完工时间"""
        self.model.setObjective(self.Cmax, GRB.MINIMIZE)

    def add_constraints(self):
        """添加所有约束条件"""
        print("======================begin 子问题1: 加工任务分配 约束=========================")
        self.add_process_assignment_constraints()  # 子问题1: 加工任务分配

        print("======================begin 子问题2: 加工任务排序 约束=========================")
        self.add_process_scheduling_constraints()  # 子问题2: 加工任务排序

        print("======================begin 子问题3: 运输任务分配 约束=========================")
        self.add_transport_assignment_constraints()  # 子问题3: 运输任务分配

        print("======================begin 子问题4: 运输任务排序 约束=========================")
        self.add_transport_scheduling_constraints()  # 子问题4: 运输任务排序

        print("======================begin 加工-运输关联 约束=========================")
        self.add_link_constraints()  # 加工-运输关联约束

        print("======================begin 目标函数 约束=========================")
        self.add_objective_constraints()  # 目标函数约束

    def add_process_assignment_constraints(self):
        """子问题1: 加工任务分配约束"""

        J = self.params['J']
        K = self.params['K']
        PM = self.params['PM']  # 加工可行性矩阵
        PP = self.params['PP']  # 加工时间矩阵

        # 公式(3): 每个工序必须分配到唯一机器
        for j in J:
            self.model.addConstr(gp.quicksum(self.phi[j, k] for k in K) == 1,
                                 name=f"process_unique_{j}")

        # 公式(4): 加工可行性约束
        for j in J:
            for k in K:
                self.model.addConstr(self.phi[j, k] <= PM[j, k],
                                     name=f"process_feasible_{j}_{k}")

        # 公式(5): 实际加工时间计算
        for j in J:
            self.model.addConstr(self.Pt[j] == gp.quicksum(self.phi[j, k] * PP[j, k] for k in K),
                                 name=f"process_time_{j}")

    def add_process_scheduling_constraints(self):
        """子问题2: 加工任务排序约束"""

        J = self.params['J']
        K = self.params['K']
        M = self.params['M']  # 大常数
        Ns = self.params['Ns']  # 相同工艺路径、同一工艺步骤的工序集合
        Nl = self.params['Nl']  # 相同工艺路径中最后一个工序

        # 公式(6)(7): 同一机器上的加工顺序约束
        for (j,j_prime,k) in self.y:
            # 获取可以在k上加工的所有工序

            # 工序j优先于j'
            self.model.addConstr(
                self.SO[j_prime] >= self.SO[j] + self.Pt[j]
                - M * (2 - self.phi[j, k] - self.phi[j_prime, k])
                - M * (1 - self.y[j, j_prime, k]),
                name=f"process_order1_{j}_{j_prime}_{k}"
            )
            # 工序j'优先于j
            self.model.addConstr(
                self.SO[j] >= self.SO[j_prime] + self.Pt[j_prime]
                - M * (2 - self.phi[j, k] - self.phi[j_prime, k])
                - M * self.y[j, j_prime, k],
                name=f"process_order2_{j}_{j_prime}_{k}"
            )

        # 公式(8): 禁止超片约束
        # group表示相同工艺路径的晶圆、相同工艺步骤的集合
        for group in Ns:
            # 确保工序按编号从小到大排序
            sorted_group = group  # 如[1,8,15,22,...]

            # 仅遍历相邻的工序对（利用传递性减少约束数量）
            for i in range(len(sorted_group) - 1):
                j_prev = sorted_group[i]  # 前一个工序（编号小）
                j_curr = sorted_group[i + 1]  # 后一个工序（编号大）

                # 遍历所有机器
                for k in K:
                    # 约束：若j_prev和j_curr在同一机器k上加工，则j_prev完成后j_curr才能开始
                    self.model.addConstr(
                        self.SO[j_curr] >= self.SO[j_prev] + self.Pt[j_prev]
                        - M * (2 - self.phi[j_prev, k] - self.phi[j_curr, k]),
                        name=f"no_overtaking_adj_{j_prev}_{j_curr}_machine_{k}"
                    )

    def add_transport_assignment_constraints(self):
        """子问题3: 运输任务分配约束"""

        J = self.params['J']
        R = self.params['R']
        TM = self.params['TM']  # 运输可行性矩阵

        # 公式9：运输任务唯一性约束
        for j in J:
            self.model.addConstr(gp.quicksum(self.mu[j, r] for r in R) == 1,
                                 name=f"transport_unique_{j}")

        # 公式(10)：运输可行性约束
        for j in J:
            for r in R:
                self.model.addConstr(self.mu[j, r] <= TM[j, r],
                                     name=f"transport_feasible_{j}_{r}")

    def add_transport_scheduling_constraints(self):
        """子问题4: 运输任务排序约束"""

        J = self.params['J']
        R = self.params['R']
        M = self.params['M']  # 大常数
        K = self.params['K']

        TT = self.params['TT'] # 运输时间
        Ns = self.params['Ns']
        Nj = self.params['Nj']

        TM = self.params['TM']

        # 添加互斥约束
        # for r in R:
        #     feasible_jobs = [j for j in J if TM.get((j, r), 0) == 1]
        #     for i in range(len(feasible_jobs)):
        #         j = feasible_jobs[i]
        #         for j_prime in feasible_jobs[i + 1:]:
        #             if (j, j_prime, r) in self.z and (j_prime, j, r) in self.z:
        #                 self.model.addConstr(
        #                     self.z[(j, j_prime, r)] + self.z[(j_prime, j, r)] == 1,
        #                     name=f"transport_mutex_{j}_{j_prime}_{r}"
        #                 )

        # 只对相邻工序添加互斥约束
        # for (j_prev, j_curr, r) in self.z:
        #     # 确保j_prev和j_curr是同一工件的相邻工序
        #     self.model.addConstr(
        #         self.z[(j_prev, j_curr, r)] + self.z.get((j_curr, j_prev, r), 0) == 1,
        #         name=f"transport_mutex_{j_prev}_{j_curr}_{r}"
        #     )


        # 仅为同一机械手、同一工件的相邻工序创建互斥约束
        # for r in R:
        #     for group in Nj:  # Nj 是同一工件的工序列表（如工件1的[1-7]）
        #         for i in range(len(group) - 1):
        #             j_prev = group[i]
        #             j_curr = group[i + 1]
        #             if TM.get((j_prev, r), 0) == 1 and TM.get((j_curr, r), 0) == 1:
        #                 self.z[(j_prev, j_curr, r)] = self.model.addVar(vtype=GRB.BINARY)
        #                 # 仅添加相邻工序的互斥约束（替代原所有工序对的循环）
        #                 self.model.addConstr(
        #                     self.z[(j_prev, j_curr, r)] + self.z.get((j_curr, j_prev, r), 0) == 1,
        #                     name=f"transport_mutex_{j_prev}_{j_curr}_{r}"
        #                 )

        # 2. 为每个机械手r，处理其可能运输的任务对
        for r in R:
            # 筛选出"可由机械手r运输"的任务集J_r（减少无关约束）
            J_r = [j for j in J if TM.get((j, r), 0) == 1]
            if len(J_r) < 2:
                continue  # 一个机械手上只分配了一个运输任务：单个任务无需排序约束


            for i in range(len(J_r)):
                # 同一个机械手上的运输任务
                j = J_r[i]
                for j_prime in J_r[i + 1:]:  # 只处理j < j_prime的对，避免重复
                    # 3.1 互斥约束：要么j先于j_prime，要么j_prime先于j（二选一）
                    self.model.addConstr(
                        self.z[(j, j_prime, r)] + self.z[(j_prime, j, r)] == 1,
                        name=f"mutex_{j}_{j_prime}_{r}"
                    )

                    # 3.2 顺序约束：当j先于j_prime时，j_prime的开始时间需晚于j完成时间
                    # 约束逻辑：若z[j,j',r]=1且两者都由r运输，则ST[j'] >= ST[j] + Tt[j]
                    self.model.addConstr(
                        self.ST[j_prime] >= self.ST[j] + self.Tt[j]
                        - M * (1 - self.z[(j, j_prime, r)])  # z=0时该条件失效
                        - M * (2 - self.mu[j, r] - self.mu[j_prime, r]),  # 任一不由r运输时失效
                        name=f"order_{j}_{j_prime}_{r}"
                    )

                    # 3.3 顺序约束：当j_prime先于j时，j的开始时间需晚于j_prime完成时间
                    self.model.addConstr(
                        self.ST[j] >= self.ST[j_prime] + self.Tt[j_prime]
                        - M * (1 - self.z[(j_prime, j, r)])  # z=0时该条件失效
                        - M * (2 - self.mu[j, r] - self.mu[j_prime, r]),  # 任一不由r运输时失效
                        name=f"order_{j_prime}_{j}_{r}"
                    )
        # 仅考虑同一个工件的相邻工序的运输顺序
        # for (j_prev, j_curr, r) in self.z:
        #     # j_prev优先于j_curr的约束
        #     self.model.addConstr(
        #         self.ST[j_curr] >= self.ST[j_prev] + self.Tt[j_prev]
        #         - M * (1 - self.z[(j_prev, j_curr, r)])
        #         - M * (2 - self.mu[j_prev, r] - self.mu[j_curr, r]),
        #         name=f"transport_order1_{j_prev}_{j_curr}_{r}"
        #     )
        #     # j_curr优先于j_prev的约束（实际不会触发，因相邻工序有固定顺序）
        #     self.model.addConstr(
        #         self.ST[j_prev] >= self.ST[j_curr] + self.Tt[j_curr]
        #         - M * self.z[(j_prev, j_curr, r)]
        #         - M * (2 - self.mu[j_prev, r] - self.mu[j_curr, r]),
        #         name=f"transport_order2_{j_prev}_{j_curr}_{r}"
        #     )
        #
        #
        #     # 公式(13)：运输禁止超片约束
        #     # group表示相同工艺路径的晶圆、相同工艺步骤的集合
        #     for group in Ns:
        #         # 确保工序按编号从小到大排序
        #         sorted_group = group  # 如[1,8,15,22,...]
        #
        #         # 仅遍历相邻的工序对（利用传递性减少约束数量）
        #         for i in range(len(sorted_group) - 1):
        #             j_prev = sorted_group[i]  # 前一个工序（编号小）
        #             j_curr = sorted_group[i + 1]  # 后一个工序（编号大）
        #
        #             # 遍历所有机器
        #             for k in K:
        #                 # 约束：若j_prev和j_curr在同一机器k上加工，则j_prev完成后j_curr才能开始
        #                 self.model.addConstr(
        #                     self.ST[j_curr] >= self.ST[j_prev] + self.Tt[j_prev],
        #                     name=f"no_overtaking_adj_{j_prev}_{j_curr}_machine_{k}"
        #                 )

    def add_link_constraints(self):
        """加工与运输关联约束（公式14-21）"""

        J = self.params['J']
        Nf = self.params['Nf']  # 首道工序集合
        alpha = self.params['alpha']  # JIT加工驻留时间（15s）
        beta = self.params['beta']    # JIT运输驻留时间（30s）
        TT = self.params['TT']        # 运输时间矩阵
        Nj = self.params['Nj']        # 每个工件的所有工序集合


        for j in J:
            # 首工序
            if j in Nf:
                # 首工序运输任务直接等于9秒
                self.model.addConstr(
                    self.Tt[j] == 9,
                    name=f"transport_time_exact_{j}"
                )
            else:
                # 中间工序运输任务直接等于10秒
                self.model.addConstr(
                    self.Tt[j] == 10,
                    name=f"transport_time_exact_{j}"
                )


        # 公式(18)：运输完成后才能开始加工
        for j in J:
            self.model.addConstr(
                self.SO[j] >= self.ST[j] + self.Tt[j],
                name=f"transport_before_process_{j}"
            )

        # 同一个工件的所有工序：后一个工序的运输开始实际必须大于前一个工序的加工完成时间
        for group in Nj:
            # group为同一工件的所有工序（按顺序排列）
            # 遍历相邻工序对（j_prev为前序工序，j_curr为后序工序）
            for i in range(len(group) - 1):
                j_prev = group[i]
                j_curr = group[i + 1]
                # 约束：j_curr的运输开始时间 > j_prev的加工完成时间
                self.model.addConstr(
                    self.ST[j_curr] >= self.SO[j_prev] + self.Pt[j_prev],
                    name=f"job_sequence_{j_prev}_{j_curr}"
                )

            # JIT驻留时间约束
            for i in range(len(group) - 1):
                j_prev = group[i]
                j_curr = group[i + 1]
                # JIT加工驻留时间约束
                self.model.addConstr(self.ST[j_curr] <= self.SO[j_prev] + self.Pt[j_prev] + alpha,
                     name=f"jit_process_residence_{j}"
                )
                # JIT运输驻留时间约束
                self.model.addConstr(self.ST[j_curr] + self.Tt[j_curr] <= self.SO[j_prev] + self.Pt[j_prev] + beta,
                    name=f"jit_transport_residence_{j}"
                )



        # JIT约束
        # for j in J:
        #         # 公式(20)：JIT加工驻留时间约束
        #         self.model.addConstr(
        #             self.ST[j] <= self.SO[j_prev] + self.Pt[j_prev],
        #             name=f"jit_process_residence_{j}"
        #         )
        #
        #         # 公式(21)：JIT运输驻留时间约束
        #         self.model.addConstr(
        #             self.ST[j] + self.Tt[j] <= self.SO[j_prev] + self.Pt[j_prev],
        #             name=f"jit_transport_residence_{j}"
        #         )

    def add_objective_constraints(self):
        """目标函数相关约束"""
        Nl = self.params['Nl']  # 最后一道工序集合

        # 公式(2): 最大完工时间约束
        for j in Nl:
            self.model.addConstr(self.Cmax >= self.SO[j] + self.Pt[j],
                                 name=f"objective_{j}")

    def optimize(self):
        """求解模型并返回结果"""

        print("====================开始求解模型=================")
        start_time = time.time()
        self.model.optimize()
        solve_time = time.time() - start_time

        # 关键：检查是否存在可行解（无论状态码如何）
        has_feasible = self.model.SolCount > 0  # SolCount：可行解数量

        if has_feasible:
            try:
                job_mapping = self.params['job_mapping']
                result = {
                    "solve_time": solve_time,
                    "status": self.model.status,  # 保留状态码（如9表示超时）
                    "Cmax": self.Cmax.X,
                    "Pt": {j: self.Pt[j].X for j in self.params['J']},
                    "Tt": {j: self.Tt[j].X for j in self.params['J']},
                    "SO": {j: self.SO[j].X for j in self.params['J']},
                    "ST": {j: self.ST[j].X for j in self.params['J']},
                    "job_id": job_mapping,
                    # 安全获取 phi（避免因浮点数误差导致的空列表）
                    "phi": {},
                    "mu": {}
                }
                # 处理 phi（机器分配）
                for j in self.params['J']:
                    machines = [k for k in self.params['K'] if abs(self.phi[j, k].X - 1) < 1e-4]
                    if not machines:
                        raise ValueError(f"工序 {j} 未分配到任何机器（phi 变量异常）")
                    result["phi"][j] = machines[0]
                # 处理 mu（机械手分配）
                for j in self.params['J']:
                    robots = [r for r in self.params['R'] if abs(self.mu[j, r].X - 1) < 1e-4]
                    if not robots:
                        raise ValueError(f"工序 {j} 未分配到任何机械手（mu 变量异常）")
                    result["mu"][j] = robots[0]
                # 补充排序变量（可选）
                result["y"] = {(j, jp, k): v.X for (j, jp, k), v in self.y.items() if v.X == 1}
                result["z"] = {(j, jp, r): v.X for (j, jp, r), v in self.z.items() if v.X == 1}
                return result
            except Exception as e:
                print(f"生成结果时出错：{e}")
                return {"solve_time": solve_time, "status": self.model.status, "error": str(e)}
        else:
            # 无可行解时的处理
            if self.model.status == GRB.INFEASIBLE:
                print("模型不可行，生成冲突约束文件...")
                self.model.computeIIS()
                self.model.write("infeasible_constraints.ilp")
                print("冲突约束已保存到 infeasible_constraints.ilp")
            print(f"求解状态: {self.model.status}（无可行解）")
            return {"solve_time": solve_time, "status": self.model.status}




if __name__ == "__main__":
    # 参数初始化
    # 任务信息
    # 创建工件
    job_creator = create_job()
    job_d = job_creator.generator_job_d()
    job_c = job_creator.generator_job_c()

    # 创建任务实例
    task = Task(
        name="任务-5c",
        job_types=[job_c, job_d],
        job_counts=[10, 40]
    )

    task.print_task_info()

    paramter_generator = ParamterGenerator(task)

    # 生成参数
    params = paramter_generator.generate_all_params()
    paramter_generator.pretty_print_dict(params)

    # # 创建并求解模型
    model = SemiconductorScheduler(params)
    result = model.optimize()

    # 5. 输出求解结果
    print_schedule_results(result)


